package com.demo.wormhole.component.rabbit.template;

import com.demo.wormhole.core.configuration.WormholeIsolationConfiguration;
import com.demo.wormhole.core.context.WormholeContext;
import com.demo.wormhole.core.context.WormholeContextHolder;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.connection.ConnectionFactory;
import org.springframework.amqp.rabbit.connection.CorrelationData;
import org.springframework.amqp.rabbit.core.RabbitTemplate;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.lang.Nullable;

public class ShadowRabbitTemplate extends RabbitTemplate {
    public ShadowRabbitTemplate(ConnectionFactory connectionFactory) {
        super(connectionFactory);
    }
    
    @Autowired
    private WormholeIsolationConfiguration isolationConfiguration;
    
    /**
     * 发送时候,需要添加头信息
     *
     * @param exchange
     * @param routingKey
     * @param message
     * @param correlationData
     * @throws AmqpException
     */
    @Override
    public void send(final String exchange, final String routingKey,
                     final Message message, @Nullable final CorrelationData correlationData)
            throws AmqpException {
        WormholeContext wormholeContext = WormholeContextHolder.getContext();
        if (null == wormholeContext) {
            super.send(exchange, routingKey, message, correlationData);
        } else {
            message.getMessageProperties()
                    .setHeader(WormholeContextHolder.WORMHOLE_REQUEST_MARK,
                            wormholeContext.toString());
            //生成Rabbit 隔离Key
            String wormholeRoutingKey = isolationConfiguration.generateIsolationKey(routingKey);
            //调用父类进行发送
            super.send(exchange, wormholeRoutingKey, message, correlationData);
        }
    }
}


